package com.mytest.example.asyncservlet.controller;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Collection;
import java.util.concurrent.*;


@Slf4j
@RequestMapping("/asyncContext")
@RestController
public class AsyncContextController {
    //定时任务，阻塞的最大超时时间。
    private ScheduledExecutorService timeoutChecker = new ScheduledThreadPoolExecutor(1, threadFactory);
    private static final ThreadFactory threadFactory = new ThreadFactoryBuilder().setNameFormat("timeout-checker-%d").build();
    private static ThreadPoolExecutor executor = new ThreadPoolExecutor(100, 200, 50000L, TimeUnit.MILLISECONDS, new ArrayBlockingQueue<>(100), threadFactory, new ThreadPoolExecutor.AbortPolicy());
    private static final Multimap<String, AsyncTask> taskCache = Multimaps.synchronizedSetMultimap(HashMultimap.create());

    public AsyncContextController() {

    }

    @GetMapping("/start")
    public void start(HttpServletRequest request, HttpServletResponse response) throws IOException {
        log.info("线程ID={},Name={},请求进来了", Thread.currentThread().getId(), Thread.currentThread().getName());
        response.setContentType("text/plain;charset=UTF-8");
        response.setDateHeader("Expires", 0);
        response.setHeader("Pragma", "no-cache");
        response.setHeader("Cache-Control", "no-cache,nos-tore");
        final PrintWriter writer = response.getWriter();
        writer.println("处理开始");
        writer.flush();
        if (request.isAsyncSupported()) {
            String taskName = request.getParameter("taskName");
            String timeOut = request.getParameter("timeOut");
            String sleep = request.getParameter("sleep");
            if (timeOut == null) {
                timeOut = "29";
            }
            if (sleep == null) {
                sleep = "5";
            }
            // 开启异步
            AsyncContext asyncContext = request.startAsync(request, response);
            asyncContext.setTimeout(0);
            AsyncTask asyncTask = new AsyncTask(asyncContext, true);
            taskCache.put(taskName, asyncTask);
            // Servlet不会被阻塞,而是直接往下执行
            String finalSleep = sleep;
            //执行自定义异步任务executeTask
            asyncContext.start(() -> {
                log.info("线程ID={},Name={},执行asyncContext.start", Thread.currentThread().getId(), Thread.currentThread().getName());

            });
            executor.execute(() -> {
                try {
                    log.info("线程ID={},Name={},执行executor.execute", Thread.currentThread().getId(), Thread.currentThread().getName());
                    TimeUnit.SECONDS.sleep(Long.parseLong(finalSleep));
                    Collection<AsyncTask> asyncTasks = taskCache.removeAll(taskName);
                    if (!CollectionUtils.isEmpty(asyncTasks)) {
                        //为每一个AsyncContext设置200的状态码以及响应数据。
                        for (AsyncTask asyncTask1 : asyncTasks) {
                            //表明未超时，已经进行处理了。
                            asyncTask1.setTimeout(false);
                            HttpServletResponse response1 = (HttpServletResponse) asyncTask1.getAsyncContext().getResponse();
                            response1.setStatus(HttpServletResponse.SC_OK);
                            final PrintWriter writer1 = response1.getWriter();
                            writer1.println("异步任务执行结束");
                            writer1.flush();
                            asyncTask1.getAsyncContext().complete();
                        }
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
            timeoutChecker.schedule(() -> {
                if (asyncTask.isTimeout()) {
                    //清除缓存中的任务
                    if (taskCache.remove(taskName, asyncTask)) {
                        response.setStatus(HttpServletResponse.SC_OK);
                        writer.println("处理超时");
                        writer.flush();
                        asyncTask.getAsyncContext().complete();
                    }
                }
            }, Integer.parseInt(timeOut), TimeUnit.SECONDS);

            writer.println("处理结束");
            writer.flush();
            log.info("线程ID={},Name={},请求出去了", Thread.currentThread().getId(), Thread.currentThread().getName());
        } else {
            writer.println("不支持异步");
            writer.flush();
            log.info("线程ID={},Name={},请求出去了", Thread.currentThread().getId(), Thread.currentThread().getName());
        }
    }

    //自定义任务对象
    @Data
    private static class AsyncTask {
        // 长轮询请求的上下文，包含请求和响应体
        private AsyncContext asyncContext;
        // 超时标记
        private boolean timeout;

        public AsyncTask(AsyncContext asyncContext, boolean timeout) {
            this.asyncContext = asyncContext;
            this.timeout = timeout;
        }
    }
}
